 /*
 * Copyright (C) 2019 Red Hat, Inc.
 *
 * Written By: Gal Hammer <ghammer@redhat.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include "viofs.h"
#include "viofs.tmh"

#ifdef ALLOC_PRAGMA
#pragma alloc_text(INIT, DriverEntry)
#pragma alloc_text(PAGE, VirtFsEvtDeviceAdd)
#pragma alloc_text(PAGE, VirtFsEvtDriverContextCleanup)
#endif

NTSTATUS DriverEntry(IN PDRIVER_OBJECT DriverObject,
                     IN PUNICODE_STRING RegistryPath)
{
    NTSTATUS status;
    WDF_DRIVER_CONFIG config;
    WDF_OBJECT_ATTRIBUTES attributes;

    ExInitializeDriverRuntime(DrvRtPoolNxOptIn);
    
    // Initialize the WPP tracing.
    WPP_INIT_TRACING(DriverObject, RegistryPath);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "--> %!FUNC!");

    // Register a cleanup callback so that we can call WPP_CLEANUP when
    // the framework driver object is deleted during driver unload.
    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.EvtCleanupCallback = VirtFsEvtDriverContextCleanup;

    WDF_DRIVER_CONFIG_INIT(&config, VirtFsEvtDeviceAdd);

    status = WdfDriverCreate(DriverObject, RegistryPath, &attributes,
        &config, WDF_NO_HANDLE);

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfDriverCreate failed: %!STATUS!", status);
        WPP_CLEANUP(DriverObject);
    }

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "<-- %!FUNC!");

    return status;
}

NTSTATUS
VirtFsAllocateIndirectPagesList(IN PDEVICE_CONTEXT Context)
{
    int i;
    VirtIODevice* dev = &Context->VDevice.VIODevice;
    RtlZeroMemory(Context->IndirectPages, sizeof(Context->IndirectPages));
    Context->UsermodeIndirectEntriesCnt = 0;

    for (i = 0; i < VIRT_FS_INDIRECT_REQUESTS_NUM; ++i)
    {
        PVIRTIO_FS_INDIRECT_PAGES Pages = &Context->IndirectPages[i];

        Pages->VirtAddr = VirtIOWdfDeviceAllocDmaMemory(dev,
            VIRT_FS_INDIRECT_PAGES * PAGE_SIZE, 0);
        if (Pages->VirtAddr == NULL)
        {
            VirtFsFreeIndirectPagesList(Context);
            return STATUS_INSUFFICIENT_RESOURCES;
        }

        Pages->PhysAddr = VirtIOWdfDeviceGetPhysicalAddress(dev, Pages->VirtAddr);

        PushEntryList(&Context->IndirectList, &Pages->ListEntry);
    }
    return STATUS_SUCCESS;
}

VOID
VirtFsFreeIndirectPagesList(IN PDEVICE_CONTEXT Context)
{
    int i;
    VirtIODevice* dev = &Context->VDevice.VIODevice;
    for (i = 0; i < VIRT_FS_INDIRECT_REQUESTS_NUM; ++i)
    {
        PVIRTIO_FS_INDIRECT_PAGES Pages = &Context->IndirectPages[i];
        if (Pages->VirtAddr != NULL)
        {
            VirtIOWdfDeviceFreeDmaMemory(dev, Pages->VirtAddr);
            Pages->VirtAddr = NULL;
        }
    }
}

NTSTATUS VirtFsEvtDeviceAdd(IN WDFDRIVER Driver,
                            IN PWDFDEVICE_INIT DeviceInit)
{
    NTSTATUS status;
    WDFDEVICE device;
    WDF_PNPPOWER_EVENT_CALLBACKS pnpPowerCallbacks;
    WDF_OBJECT_ATTRIBUTES attributes;
    WDFQUEUE queue;
    WDF_IO_QUEUE_CONFIG queueConfig;
    WDF_INTERRUPT_CONFIG interruptConfig;
    PDEVICE_CONTEXT context;

    UNREFERENCED_PARAMETER(Driver);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "--> %!FUNC!");

    PAGED_CODE();

    WDF_PNPPOWER_EVENT_CALLBACKS_INIT(&pnpPowerCallbacks);

    pnpPowerCallbacks.EvtDevicePrepareHardware = VirtFsEvtDevicePrepareHardware;
    pnpPowerCallbacks.EvtDeviceReleaseHardware = VirtFsEvtDeviceReleaseHardware;
    pnpPowerCallbacks.EvtDeviceD0Entry = VirtFsEvtDeviceD0Entry;
    pnpPowerCallbacks.EvtDeviceD0Exit = VirtFsEvtDeviceD0Exit;

    WdfDeviceInitSetPnpPowerEventCallbacks(DeviceInit, &pnpPowerCallbacks);
    WdfDeviceInitSetIoType(DeviceInit, WdfDeviceIoDirect);

    WDF_OBJECT_ATTRIBUTES_INIT_CONTEXT_TYPE(&attributes, DEVICE_CONTEXT);
    attributes.EvtCleanupCallback = VirtFsEvtDeviceContextCleanup;

    status = WdfDeviceCreate(&DeviceInit, &attributes, &device);
    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfDeviceCreate failed: %!STATUS!", status);
        return status;
    }

    context = GetDeviceContext(device);

    WDF_INTERRUPT_CONFIG_INIT(&interruptConfig,
        VirtFsEvtInterruptIsr, VirtFsEvtInterruptDpc);

    interruptConfig.EvtInterruptEnable = VirtFsEvtInterruptEnable;
    interruptConfig.EvtInterruptDisable = VirtFsEvtInterruptDisable;

    for (ULONG idx = 0; idx < ARRAYSIZE(context->WdfInterrupt); idx++)
    {
        status = WdfInterruptCreate(device, &interruptConfig,
            WDF_NO_OBJECT_ATTRIBUTES, &context->WdfInterrupt[idx]);

        if (!NT_SUCCESS(status))
        {
            TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
                "WdfInterruptCreate #%ul failed: %!STATUS!", idx, status);
            return status;
        }
    }

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.ParentObject = device;
    status = WdfSpinLockCreate(&attributes,
        &context->RequestsLock);
    if (NT_SUCCESS(status))
    {
        status = WdfSpinLockCreate(&attributes,
            &context->IndirectListLock);
    }

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfSpinLockCreate failed: %!STATUS!", status);
        return status;
    }

    status = WdfDeviceCreateDeviceInterface(device,
        &GUID_DEVINTERFACE_VIRT_FS, NULL);

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfDeviceCreateDeviceInterface failed: %!STATUS!", status);
        return status;
    }

    WDF_IO_QUEUE_CONFIG_INIT_DEFAULT_QUEUE(&queueConfig, WdfIoQueueDispatchSequential);
    queueConfig.EvtIoDeviceControl = VirtFsEvtIoDeviceControl;
    queueConfig.EvtIoStop = VirtFsEvtIoStop;
    queueConfig.AllowZeroLengthRequests = FALSE;

    status = WdfIoQueueCreate(device, &queueConfig, WDF_NO_OBJECT_ATTRIBUTES,
        &queue);

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfIoQueueCreate failed: %!STATUS!", status);
        return status;
    }

    WDF_OBJECT_ATTRIBUTES_INIT(&attributes);
    attributes.ParentObject = device;

    status = WdfLookasideListCreate(&attributes, sizeof(VIRTIO_FS_REQUEST),
        NonPagedPool, WDF_NO_OBJECT_ATTRIBUTES, VIRT_FS_MEMORY_TAG,
        &context->RequestsLookaside);

    if (!NT_SUCCESS(status))
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_INIT,
            "WdfLookasideListCreate failed: %!STATUS!", status);
        return status;
    }

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "<-- %!FUNC!");

    return status;
}

VOID VirtFsEvtDeviceContextCleanup(IN WDFOBJECT DeviceObject)
{
    PDEVICE_CONTEXT context = GetDeviceContext(DeviceObject);
    PSINGLE_LIST_ENTRY iter;

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "--> %!FUNC!");

    WdfSpinLockAcquire(context->RequestsLock);
    iter = PopEntryList(&context->RequestsList);
    while (iter != NULL)
    {
        PVIRTIO_FS_REQUEST fs_req = CONTAINING_RECORD(iter,
            VIRTIO_FS_REQUEST, ListEntry);

        FreeVirtFsRequest(context, fs_req);

        iter = PopEntryList(&context->RequestsList);
    };
    WdfSpinLockRelease(context->RequestsLock);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "<-- %!FUNC!");
}

VOID VirtFsEvtDriverContextCleanup(IN WDFOBJECT DriverObject)
{
    UNREFERENCED_PARAMETER(DriverObject);

    TraceEvents(TRACE_LEVEL_VERBOSE, DBG_INIT, "<-> %!FUNC!");

    PAGED_CODE();

    // Stop the WPP tracing.
    WPP_CLEANUP(WdfDriverWdmGetDriverObject(DriverObject));
}

void FreeVirtFsRequest(IN PDEVICE_CONTEXT Context,
                       IN PVIRTIO_FS_REQUEST fs_req)
{
    if (fs_req->IndirectPages)
    {
        FreeIndirectPages(Context, fs_req->IndirectPages);
        fs_req->IndirectPages = NULL;
    }

    if (fs_req->RequestType == VIRTIO_FS_REQUEST_TYPE_WDF)
    {
        VIRTIO_FS_REQUEST_WDF* Request = &fs_req->u.Wdf;
        if (Request->InputBuffer != NULL)
        {
            MmFreePagesFromMdl(Request->InputBuffer);
            ExFreePool(Request->InputBuffer);
            Request->InputBuffer = NULL;
            Request->InputBufferLength = 0;
        }

        if (Request->OutputBuffer != NULL)
        {
            MmFreePagesFromMdl(Request->OutputBuffer);
            ExFreePool(Request->OutputBuffer);
            Request->OutputBuffer = NULL;
            Request->OutputBufferLength = 0;
        }
    }

    if (fs_req->Handle != NULL)
    {
        WdfObjectDelete(fs_req->Handle);
        fs_req->Handle = NULL;
    }
}

static BOOLEAN AllocateIndirectEntry(PVIRTIO_FS_INDIRECT_PAGES *ppIndirectPages, ULONG SgCount)
{
    PHYSICAL_ADDRESS LowAddress;
    PHYSICAL_ADDRESS HighAddress;
    PHYSICAL_ADDRESS SkipAddress;
    PVIRTIO_FS_INDIRECT_PAGES pPages;
    ULONG  pagesCnt;

    *ppIndirectPages = NULL;
    pPages = ExAllocatePoolUninitialized(NonPagedPool, sizeof(*pPages), VIRT_FS_MEMORY_TAG);
    if (!pPages)
    {
        return FALSE;
    }

    pagesCnt = (SgCount + VIRT_FS_INDIRECT_PAGE_CAPACITY - 1) / VIRT_FS_INDIRECT_PAGE_CAPACITY;

    LowAddress.QuadPart = 0;
    HighAddress.QuadPart = -1; // no limit
    SkipAddress.QuadPart = 0;

    pPages->AllocatedMdl = MmAllocatePagesForMdlEx(
        LowAddress, HighAddress, SkipAddress, pagesCnt * PAGE_SIZE,
        MmCached, MM_ALLOCATE_REQUIRE_CONTIGUOUS_CHUNKS);
    if (!pPages->AllocatedMdl)
    {
        ExFreePool(pPages);
        return FALSE;
    }

    pPages->VirtAddr = MmMapLockedPagesSpecifyCache(
        pPages->AllocatedMdl, KernelMode,
        MmCached, NULL, FALSE, NormalPagePriority | MdlMappingNoExecute);
    if (!pPages->VirtAddr)
    {
        MmFreePagesFromMdl(pPages->AllocatedMdl);
        ExFreePool(pPages->AllocatedMdl);
        ExFreePool(pPages);
        return FALSE;
    }

    PPFN_NUMBER pfnArray = MmGetMdlPfnArray(pPages->AllocatedMdl);
    pPages->PhysAddr.QuadPart = pfnArray[0] << PAGE_SHIFT;

    *ppIndirectPages = pPages;

    return TRUE;
}

static VOID FreeIndirectEntry(PVIRTIO_FS_INDIRECT_PAGES pPages)
{
    MmUnmapLockedPages(pPages->VirtAddr, pPages->AllocatedMdl);
    MmFreePagesFromMdl(pPages->AllocatedMdl);
    ExFreePool(pPages->AllocatedMdl);
    ExFreePool(pPages);
}

NTSTATUS GetIndirectPages(
    PDEVICE_CONTEXT Context, ULONG SgCount, BOOLEAN IsUsermode,
    PVIRTIO_FS_INDIRECT_PAGES *ppIndirectPages)
{
    PSINGLE_LIST_ENTRY iter;

    if (SgCount > VIRT_FS_INDIRECT_CAPACITY)
    {
        TraceEvents(TRACE_LEVEL_ERROR, DBG_IOCTL,
            "The buffer is too big to fit into vio-queue: buf_count %u", SgCount);
        return STATUS_INVALID_PARAMETER;
    }

    if (IsUsermode &&
        InterlockedIncrement(&Context->UsermodeIndirectEntriesCnt) > (VIRT_FS_INDIRECT_REQUESTS_NUM / 2))
    {
        InterlockedDecrement(&Context->UsermodeIndirectEntriesCnt);
        TraceEvents(TRACE_LEVEL_ERROR, DBG_IOCTL,
            "Only one IoRequest from usermode is allowed at once");
        return STATUS_ACCESS_DENIED;
    }

    WdfSpinLockAcquire(Context->IndirectListLock);
    iter = PopEntryList(&Context->IndirectList);
    WdfSpinLockRelease(Context->IndirectListLock);
    if (iter == NULL)
    {
        TraceEvents(TRACE_LEVEL_WARNING, DBG_IOCTL, "Run out of preallocated indirect pages. Allocating from pool.");
        if (!AllocateIndirectEntry(ppIndirectPages, SgCount))
        {
            if (IsUsermode)
            {
                InterlockedDecrement(&Context->UsermodeIndirectEntriesCnt);
            }
            return STATUS_NO_MEMORY;
        }
    }
    else
    {
        *ppIndirectPages = CONTAINING_RECORD(iter, VIRTIO_FS_INDIRECT_PAGES, ListEntry);
    }

    (*ppIndirectPages)->IsUsermode = IsUsermode;

    return STATUS_SUCCESS;
}

VOID FreeIndirectPages(PDEVICE_CONTEXT Context, PVIRTIO_FS_INDIRECT_PAGES Pages)
{
    if (Pages == NULL)
    {
        return;
    }

    if (Pages->IsUsermode)
    {
        InterlockedDecrement(&Context->UsermodeIndirectEntriesCnt);
    }

    if (Pages->AllocatedMdl)
    {
        // not a preallocated IndirectEntry
        FreeIndirectEntry(Pages);
        return;
    }

    WdfSpinLockAcquire(Context->IndirectListLock);
    PushEntryList(&Context->IndirectList, &Pages->ListEntry);
    WdfSpinLockRelease(Context->IndirectListLock);
}
